from fastapi import APIRouter, Query
from pydantic import BaseModel

from service.rag_service import rag_query_by_file, qa
from utils.base_response import BaseResponse

rag = APIRouter()

# 添加请求模型
class QueryRequest(BaseModel):
    user_input: str

@rag.post("/get")
async def rag_query(request: QueryRequest):  # 修改为接收请求体
    """
    使用RAG模型进行映射规约查询
    :param request: 包含user_input的请求体
    :return:
    """
    result = await rag_query_by_file(request.user_input)
    return BaseResponse(code=200, msg="success", data=result)

@rag.post("/get_all")
async def rag_query(request: QueryRequest):  # 修改为接收请求体
    """
    使用RAG模型进行全量搜索查询
    :param request: 包含user_input的请求体
    :return:
    """
    result = await qa(request.user_input)
    return BaseResponse(code=200, msg="success", data=result)

